package com.zyk.plugin;

import com.zyk.Environment;
import com.zyk.common.PackageScanner;
import com.zyk.utils.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.*;

public class DispatcherTransformer implements ClassFileTransformer {
    Logger LOGGER = LoggerFactory.getLogger(DispatcherTransformer.class);

    Set<String> hookedClassSet = new HashSet<>();
    Map<String, List<JaTransformer>> transformerMap = new HashMap<>();

    public void addTransformer(JaTransformer transformer) {
        Environment environment = Environment.getInstance();
        if (null == transformer) {
            return;
        }

        if (environment.isAttachMode() && (transformer.mode() & JaTransformer.ATTACH_MODE) == 0) {
            LOGGER.debug("Transformer: {} is set to not load in attach mode, ignored.", transformer.getClass().getName());
            return;
        }

        if (environment.isJavaagentMode() && (transformer.mode() & JaTransformer.AGENT_MODE) == 0) {
            LOGGER.debug("Transformer: {} is set to not load in agent mode, ignored.", transformer.getClass().getName());
            return;
        }

        synchronized (this) {
            String packageName = transformer.getHookPackageName();
            if (!StringUtils.isEmpty(packageName)) {
                for (String className : PackageScanner.scan(packageName)) {
                    hookedClassSet.add(className);
                    transformerMap.computeIfAbsent(className, k -> new ArrayList<>()).add(transformer);
                }
            }

            String className = transformer.getHookClassName();
            if (!StringUtils.isEmpty(className)) {
                hookedClassSet.add(className);
                transformerMap.computeIfAbsent(className, k -> new ArrayList<>()).add(transformer);
            }
        }
    }

    public void addTransformers(List<JaTransformer> transformers) {
        if (null == transformers) {
            return;
        }

        for (JaTransformer transformer : transformers) {
            addTransformer(transformer);
        }
    }

    public Set<String> getHookClassNames() {
        return hookedClassSet;
    }


    @Override
    public byte[] transform(ClassLoader loader, String className,
                            Class<?> classBeingRedefined, ProtectionDomain protectionDomain,
                            byte[] classfileBuffer) throws IllegalClassFormatException {
        if (null == className) {
            return null;
        }
        String loadOrRedefineClassName = className.replace('/', '.');
        if (!hookedClassSet.contains(loadOrRedefineClassName)) {
            return null;
        }

        List<JaTransformer> transformers = transformerMap.get(loadOrRedefineClassName);
        if (null != transformers) {
            for (JaTransformer transformer : transformers) {
                classfileBuffer = classBeingRedefined == null ?
                        transformer.loadClass(loader, loadOrRedefineClassName, protectionDomain, classfileBuffer)
                        :
                        transformer.redefineClass(loader, loadOrRedefineClassName, classBeingRedefined, protectionDomain, classfileBuffer);
            }
            hookedClassSet.remove(loadOrRedefineClassName);
            transformerMap.remove(loadOrRedefineClassName);
            return classfileBuffer;
        }
        return null;
    }
}
